# 处理数据

import torch
import numpy as np
from template import *
from datasets import load_dataset
from functools import partial
from pet_config import *

def convert_example(example: dict,
                    tokenizer,
                    max_seq_len: int,
                    max_label_len: int,
                    hard_template: HardTemplate,
                    train_mode=True,
                    return_tensor = False) -> dict:
    tokenized_output = {
        'input_ids':[],
        'token_type_ids': [],
        'attention_mask': [],
        'mask_positions': [],
        'mask_labels': []
    }

    for i, example in enumerate(example['text']):
        if train_mode:

            label, content = example.strip().split('\t')
        else:

            content = example.strip()

        inputs_dict = {
            'textA': content,
            'MASK': '[MASK]'
        }

        encoded_inputs = hard_template(
            inputs_dict=inputs_dict,
            tokenizer=tokenizer,
            max_seq_len=max_seq_len,
            mask_length=max_label_len
        )

        tokenized_output['input_ids'].append(encoded_inputs['input_ids'])
        tokenized_output['token_type_ids'].append(encoded_inputs['token_type_ids'])
        tokenized_output['attention_mask'].append(encoded_inputs['attention_mask'])
        tokenized_output['mask_positions'].append(encoded_inputs['mask_position'])

        if train_mode:
            label_encoded = tokenizer(text=[label])
            label_encoded = label_encoded['input_ids'][0][1: -1]
            label_encoded = label_encoded[:max_label_len]
            label_encoded = label_encoded + [tokenizer.pad_token_id] * (max_label_len - len(label_encoded))
            tokenized_output['mask_labels'].append(label_encoded)

    for k, v in tokenized_output.items():
        if return_tensor:
            tokenized_output[k] = torch.LongTensor(v)

        else:
            tokenized_output[k] = np.array(v)

    return tokenized_output









if __name__ == '__main__':
    pc = ProjectConfig()
    train_dataset = load_dataset('text', data_files=pc.train_path)
    # print(train_dataset)

    tokenizer = AutoTokenizer.from_pretrained(pc.pre_model)
    hard_template = HardTemplate(prompt='这是一条{MASK}评论：{textA}')

    convert_func = partial(convert_example,
                           tokenizer=tokenizer,
                           hard_template=hard_template,
                           max_seq_len=30,
                           max_label_len= 2)

    dataset = train_dataset.map(convert_func, batched=True)

    print(f'dataset-->{dataset}')



